# %% [markdown]
# # [基于序列信息的microRNA和gene关系对预测](https://datafountain.cn/competitions/534)

# %%
import pandas as pd
from models import BaseInfo, BaseSeq2Vec, RelationShipModel
import torch
import numpy as np
from tqdm import tqdm
import os
import sys
import gc
# GCNFrame依赖于torch1.7.*
PRETRAINED_PATH = './user_data/'
OUTPUT_PATH = './user_data/'
DATA_PATH = './datasets/'
print(sys.argv)
if len(sys.argv) == 4:
    DATA_PATH = sys.argv[1]
    OUTPUT_PATH = sys.argv[2]
    PRETRAINED_PATH = sys.argv[3]

# %% [markdown]
# ## 进行简单的预处理

# %%
dataset = pd.read_csv(f'{DATA_PATH}train_dataset/Train.csv')
mirna_seqdf = pd.read_csv(f'{DATA_PATH}train_dataset/mirna_seq.csv', index_col=0)
gene_seqdf = pd.read_csv(f'{DATA_PATH}train_dataset/gene_seq.csv', index_col=0)
predict_set = pd.read_csv(f'{DATA_PATH}test_dataset.csv')# gene,miRNA

# 清除换行符
mirna_seqdf['seq'] = mirna_seqdf['seq'].apply(lambda x:x.strip())

# %% [markdown]
# ## 加载已经训练好的特征提取器

# %%
if os.path.exists(f'{OUTPUT_PATH}base_info.m'):
    base_info = BaseInfo.load(f'{OUTPUT_PATH}base_info.m')
else:
    base_info = BaseInfo()

if os.path.exists(f'{PRETRAINED_PATH}rna2vec.m'):
    rna2vec = BaseSeq2Vec.load(f'{PRETRAINED_PATH}rna2vec.m').eval()
else:
    rna2vec = BaseSeq2Vec(base_info.get_RNAf().shape[-1], 32, base_info.get_rna_s2i(), base_info.get_RNAf(), base_info.k).eval()

if os.path.exists(f'{PRETRAINED_PATH}dna2vec.m'):
    dna2vec = BaseSeq2Vec.load(f'{PRETRAINED_PATH}dna2vec.m').eval()
else:
    dna2vec = BaseSeq2Vec(base_info.get_DNAf().shape[-1], 32, base_info.get_dna_s2i(), base_info.get_DNAf(), base_info.k).eval()

# %% [markdown]
# ## 提取训练集特征
# 参照`Train.csv`中DNA与RNA的对应关系，将序列名字换成特征向量，最后得到两个可以相互对应的特征矩阵

# %%
def base2feature(base_name:list, seq:pd.Series, transform:BaseSeq2Vec, gaps:list=[0]) -> torch.Tensor:
    '''
    base_name:碱基序列名字的列表，输出按照顺序组织特征向量
    seq:碱基序列名字与对应的序列(名字作为index)
    transform:可以将序列转换成特征向量的转换器
    '''
    if not transform.trained:
        print('Warning: Using untrained BaseSeq2Vec')
    # 获取需要转换的序列名字的列表
    name_list = list(set(base_name))
    # 获取其对应的特征
    name2vec = dict()
    with torch.no_grad():
        for name in tqdm(name_list):
            tmp = []
            for gap in gaps:
                tmp.append(
                    transform(seq[name], gap).unsqueeze(0)
                )
            name2vec[name] = torch.concat(tmp, dim=0).unsqueeze(0)
    # 按照base_name的顺序组织输出
    ans = []
    for name in tqdm(base_name):
        ans.append(name2vec[name])
    return torch.concat(ans, dim=0)

# %%
# 获取训练集的RNA特征矩阵
print('preparing train_rna_feature')
if not os.path.exists(f'{OUTPUT_PATH}train_rna_feature.tsr'):
    train_rna_feature = base2feature(dataset['miRNA'].tolist(), mirna_seqdf['seq'], rna2vec, [0, 1, 2])
    torch.save(train_rna_feature, f'{OUTPUT_PATH}train_rna_feature.tsr')
gc.collect()

# 获取预测集的RNA特征矩阵
print('preparing test_rna_feature')
if not os.path.exists(f'{OUTPUT_PATH}test_rna_feature.tsr'):
    test_rna_feature = base2feature(predict_set['miRNA'].tolist(), mirna_seqdf['seq'], rna2vec, [0, 1, 2])
    torch.save(test_rna_feature, f'{OUTPUT_PATH}test_rna_feature.tsr')
gc.collect()

# 获取训练集的DNA特征矩阵
print('preparing train_dna_feature')
if not os.path.exists(f'{OUTPUT_PATH}train_dna_feature.tsr'):
    train_dna_feature = base2feature(dataset['gene'].tolist(), gene_seqdf['sequence'], dna2vec, [0, 1, 2])
    torch.save(train_dna_feature, f'{OUTPUT_PATH}train_dna_feature.tsr')
gc.collect()

# 获取预测集的DNA特征矩阵
print('preparing test_dna_feature')
if not os.path.exists(f'{OUTPUT_PATH}test_dna_feature.tsr'):
    test_dna_feature = base2feature(predict_set['gene'].tolist(), gene_seqdf['sequence'], dna2vec, [0, 1, 2])
    torch.save(test_dna_feature, f'{OUTPUT_PATH}test_dna_feature.tsr')
gc.collect()

# %% [markdown]
# ## 获取训练集标签

# %%
print('preparing labels')
if not os.path.exists(f'{OUTPUT_PATH}labels.tsr'):
    label_dict = {
        'Functional MTI': 1.,
        'Non-Functional MTI': 0.}
    labels = dataset['label'].map(lambda x: label_dict[x]).values
    labels = torch.tensor(labels, dtype=torch.float32)
    torch.save(labels, f'{OUTPUT_PATH}labels.tsr')

# %% [markdown]
# ## 搭建用于预测的模型

# %% [markdown]
# ## 训练模型

# %%
if not os.path.exists(f'{OUTPUT_PATH}model.pth'):
    train_dna_feature = torch.load(f'{OUTPUT_PATH}train_dna_feature.tsr')
    train_rna_feature = torch.load(f'{OUTPUT_PATH}train_rna_feature.tsr')
    labels = torch.load(f'{OUTPUT_PATH}labels.tsr')
    
    model = RelationShipModel(train_rna_feature.shape[-1], train_rna_feature.shape[-2])
    model.fit(train_dna_feature, train_rna_feature, labels, 20, 150, 1e-3)
    torch.save(model, f'{OUTPUT_PATH}model.pth')
    
    del train_dna_feature, train_rna_feature, labels
    gc.collect()
model = torch.load(f'{OUTPUT_PATH}model.pth')

# %% [markdown]
# ## 开始预测

# %%
test_dna_feature = torch.load(f'{OUTPUT_PATH}test_dna_feature.tsr')
test_rna_feature = torch.load(f'{OUTPUT_PATH}test_rna_feature.tsr')

ans = model.predict(test_dna_feature, test_rna_feature, 150)
predict_set['ans'] = ans
predict_set['results'] = predict_set['ans'].map(lambda x:int(x+0.5))
predict_set.to_csv(f'{OUTPUT_PATH}ans.csv', sep=',', index=False)


